import torch
from torch.distributions.categorical import Categorical
import numpy as np
from collections import namedtuple
from tqdm import tqdm
from pyecharts.charts import Line
# from env import Minesweeper
from Controller import Director
from ppo import PPO
import time
from ppoEnv import PPODirector

Transition = namedtuple('Transition', ['state', 'ac', 'ac_prob', 'reward', 'done'])

class PPOModel:
    def __init__(self, set, batch_size=32, a_lr=0.0001, b_lr=0.002, gama=0.995, epsilon=0.2, up_time=10, epoch=50):
        self.set = set
        self.batch_size = batch_size
        self.a_lr = a_lr
        self.b_lr = b_lr
        self.gama = gama
        self.epsilon = epsilon
        self.up_time = up_time
        self.epoch = epoch

    def getAction(self, a):
        rem = max(self.set.width, self.set.height)
        x = a // rem
        y = a // rem
        return [x, y]

    def train(self, times):
        start_flag = False
        env = PPODirector(self.set)
        net = PPO([self.set.height, self.set.width], self.up_time, self.batch_size, self.a_lr, self.b_lr,
                  self.gama, self.epsilon)

        Result = []

        for i in range(times):
            with tqdm(total=self.epoch, desc='Iteration %d' % i) as pbar:
                for e in range(self.epoch):

                    s = torch.tensor(env.getStatus(), dtype=torch.float32)
                    start_flag = False
                    while not env.isEnd() and env.t < 51:
                        a,a_p = net.get_action(s)
                        action = self.getAction(a[0])
                        if not start_flag:
                            env.reset(action)
                            start_flag = True
                        [s_t,r,d] = env.updateAction(action)
                        buffer = Transition(s, a, a_p, r, d)
                        net.appdend(buffer)
                        s = s_t

                    R = np.array(env.R).sum()
                    Result.append(R)
                    if len(net.suffer) > self.batch_size:
                        net.update()
                    pbar.set_postfix({'return': '%.2f' % R})
                    pbar.update(1)

        torch.save(net.action, 'net_model.pt')
        Re = []
        for i in range(int(len(Result) / self.epoch)):
            idx = i * self.epoch
            Re.append(sum(Result[idx:idx + self.epoch]) / self.epoch)
        x = [str(i) for i in range(len(Re))]
        line = Line()
        line.add_xaxis(xaxis_data=x)
        line.add_yaxis(y_axis=Re, series_name='Recall')
        line.render('result.html')
